In [1]:
import torch
from matplotlib import pyplot as plt
from matplotlib import animation
import numpy as np
In [2]:
forecast_folder = "./results/full-dataset/forecasts/stconvs2s-r/"
In [3]:
target = torch.load(forecast_folder + 'shallow_water_step20_4_20201221-062618_target.pt')
output = torch.load(forecast_folder + 'shallow_water_step20_4_20201221-062618_output.pt')
In [4]:
np_target = target.cpu().numpy()
np_output = output.cpu().numpy()

Seq 5, steps 10 -> Seq 20, steps 10

Distance between each step: 15 min (total: 5h)

In [5]:
%matplotlib auto

fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(20,10))
ax1.title.set_text('Target')
ax1.title.set_fontsize(40)
ax2.title.set_text('Output')
ax2.title.set_fontsize(40)
data_target = np_target[0,0,0,:,:]
data_output = np_output[0,0,0,:,:]
im1 = ax1.imshow(data_target)
im2 = ax2.imshow(data_output)

def init():
    im1.set_data(np_target[0,0,0,:,:])
    im2.set_data(np_output[0,0,0,:,:])

def animate(i):
    im1.set_data(np_target[0,0,i,:,:])
    im2.set_data(np_output[0,0,i,:,:])
    
frames = np_target.shape[2]

anim = animation.FuncAnimation(fig, animate, init_func=init, frames=frames, repeat = True)

from IPython.display import HTML
HTML(anim.to_jshtml())

#anim.to_html5_video()

Using matplotlib backend: agg
Out[5]:


Once Loop Reflect
In [7]:
fig2, (ax21, ax22) = plt.subplots(1, 2, figsize=(20,10))
ax21.title.set_text('Target')
ax21.title.set_fontsize(40)
ax22.title.set_text('Output')
ax22.title.set_fontsize(40)
data_target = np_target[0,0,0,:,:]
data_output = np_output[0,0,0,:,:]
im21 = ax21.imshow(data_target)
im22 = ax22.imshow(data_output)

def init():
    im21.set_data(np.zeros((40, 40)))
    im22.set_data(np.zeros((40, 40)))

def animate(i):
    im21.set_data(np_target[i*5,0,19,:,:])
    im22.set_data(np_output[i*5,0,19,:,:])
    
frames = 200

anim2 = animation.FuncAnimation(fig2, animate, init_func=init, frames=frames, repeat = True)

from IPython.display import HTML
HTML(anim2.to_jshtml())
Out[7]:


Once Loop Reflect